import pymysql
import pymysql
from pymysql import Error
from fastapi import FastAPI, HTTPException, Query, UploadFile, Body
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from datetime import datetime,timedelta
import random
import sys
import asyncio
from pathlib import Path


# 从servers模块导入所需功能
from servers.emailcode import verification_codes,send_email,cleanup_expired_codes
from servers.mysqlconnect import create_connection
from contextlib import asynccontextmanager
from servers.LLM import llm,llm_stream

# 检查必要的依赖是否已安装
def check_dependencies():
    """检查必要的依赖是否已安装"""
    try:
        # 导入主要依赖以验证安装
        import fastapi
        import pymysql
        import uvicorn
        
        return True
    except ImportError as e:
        missing_package = str(e).split()[3].strip("'")
        return False

# 确保static目录存在
def ensure_static_dir():
    static_dir = Path("static")
    if not static_dir.exists():
        static_dir.mkdir(exist_ok=True)

# 启动应用时开始清理任务
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 应用启动时
    asyncio.create_task(cleanup_expired_codes())
    yield
def generate_code():
    """生成 6 位随机验证码"""
    return str(random.randint(100000, 999999))


# FastAPI创建应用程序
app = FastAPI(title="My API", version="1.0.0",lifespan=lifespan)

# 配置 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应指定具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 托管静态目录
app.mount("/static", StaticFiles(directory="static"), name="static")


# 注册 - 保持与客户端API一致的路由名称
@app.post("/sinup")
def sinup(email: str, pwd: str, code: str, user_name: str = None):
    connection = create_connection()
    try:
        # 检查连接是否有效
        if connection is None or not connection.open:
            connection = create_connection()
            if connection is None:
                raise Exception("无法连接到数据库")
        
        cursor = connection.cursor()
        
        # 1. 验证验证码
        if email not in verification_codes or verification_codes[email]['code'] != code:
            return {"success": False, "message": "验证码错误"}
        
        # 检查验证码是否过期
        if verification_codes[email]['expires_at'] < datetime.now():
            return {"success": False, "message": "验证码已过期，请重新获取"}
        
        # 2. 检查邮箱是否已注册
        cursor.execute("SELECT * FROM user WHERE email=%s", (email,))
        existing_user = cursor.fetchone()
        if existing_user:
            return {"success": False, "message": "该邮箱已被注册"}
        
        # 3. 处理用户名
        # 如果提供了自定义用户名，使用它；否则使用邮箱前缀
        if user_name is not None and str(user_name).strip():
            user_name = str(user_name).strip()
            
            # 检查用户名是否已存在
            cursor.execute("SELECT * FROM user WHERE user_name=%s", (user_name,))
            existing_user_name = cursor.fetchone()
            if existing_user_name:
                return {"success": False, "message": "用户名已被使用，请更换一个"}
        else:
            # 使用邮箱前缀作为用户名
            email_prefix = email.split('@')[0] if email.split('@')[0] else 'default_user'
            user_name = email_prefix
            
            # 确保用户名唯一性
            cursor.execute("SELECT * FROM user WHERE user_name=%s", (user_name,))
            counter = 1
            while cursor.fetchone():
                user_name = f"{email_prefix}_{counter}"
                cursor.execute("SELECT * FROM user WHERE user_name=%s", (user_name,))
                counter += 1
        
        # 4. 创建用户
        cursor.execute(
            "INSERT INTO user (email, pwd, user_name) VALUES (%s, %s, %s)",
            (email, pwd, user_name)
        )
        connection.commit()
        
        # 获取新创建的用户ID
        user_id = cursor.lastrowid
        
        # 5. 删除已使用的验证码
        del verification_codes[email]
        
        return {"success": True, "user": {"user_id": user_id, "email": email, "user_name": user_name}}
        
    except Exception as e:
        # 尝试回滚事务
        if connection and connection.open:
            try:
                connection.rollback()
            except:
                pass
        
        # 错误处理
        if "用户名已被使用" in str(e):
            return {"success": False, "message": "用户名已被使用，请更换一个"}
        elif "邮箱已被注册" in str(e):
            return {"success": False, "message": "该邮箱已被注册"}
        else:
            return {"success": False, "message": "注册失败，请稍后重试"}
        
        
# 登录
class LoginRequest(BaseModel):
    account: str
    pwd: str

@app.post("/login")
def login(request: LoginRequest = Body(...)):
    connection = create_connection()
    try:
        # 检查连接是否有效，如果无效则重新创建连接
        try:
            # 尝试执行简单查询来测试连接
            if connection is None or not connection.open:
                connection = create_connection()
            else:
                cursor = connection.cursor()
                cursor.execute("SELECT 1")
        except Exception as conn_error:
            connection = create_connection()
            if connection is None:
                raise Exception("无法重新连接到数据库")
                
        cursor = connection.cursor()
        # 首先尝试通过用户名查找用户
        cursor.execute("SELECT * FROM user WHERE user_name = %s", (request.account,))
        result = cursor.fetchone()
        
        # 如果用户名不存在，再尝试通过邮箱查找
        if not result:
            cursor.execute("SELECT * FROM user WHERE email = %s", (request.account,))
            result = cursor.fetchone()
        
        if result:
            if result['pwd']==request.pwd:
                return {"success": True, "user":{"user_id":result['user_id'],"email":result['email'],"user_name":result['user_name']}}
            else:
                return {"success": False, "message": "密码错误"}
        else:
            return {"success": False, "message": "账户不存在"}
    except Exception as e:
        # 如果是连接错误，尝试重新连接
        if isinstance(e, (pymysql.OperationalError, pymysql.ProgrammingError)):
            pass
        return {"success": False, "message": "服务器内部错误，请稍后重试"}


class getemail(BaseModel):
    email: str
# 发送验证码
@app.post("/send_Ver_code")
def send_Ver_code(data: getemail):
    email = data.email
    code = generate_code()
    expires_at = datetime.now() + timedelta(minutes=5)
    if send_email(email, code):
        # 存储验证码信息
        verification_codes[email] = {
            'code': code,
            'expires_at': expires_at,
            'created_at': datetime.now(),
            'attempts': 0  # 验证尝试次数
        }
        return {"success": True, "message": "验证码已发送"}
    else:
        return {"success": False, "message": "验证码发送失败"}
    
    
class changepwd(BaseModel):
    email: str
    code: str
    newpwd: str

# 修改密码
@app.post("/pwd_change")
def pwd_change(data: changepwd):
    email = data.email
    code = data.code
    newpwd = data.newpwd
    try: 
    
        if verification_codes[email]['expires_at'] < datetime.now():
            del verification_codes[email]  # 删除过期验证码
            return {"success": False, "message": "验证码已过期，请重新获取"}
        elif verification_codes[email]['code'] == code:
            connection = create_connection()
            cursor = connection.cursor()
            cursor.execute("SELECT * FROM user WHERE email=%s", (email,))
            user = cursor.fetchone()
            
            if not user:
                return {"success": False, "message": "用户不存在"}
            else:
                cursor.execute("UPDATE user SET pwd=%s WHERE email=%s", (newpwd, email))
                connection.commit()
                return {"success": True, "message": "密码修改成功"}
        else:
            return {"success": False, "message": "验证码错误"}
    except Exception as e:
        return {"success": False, "message": "服务器内部错误，请稍后重试"}

class changeuname(BaseModel):
    user_id: int
    user_name: str

    
# 修改用户名
@app.post("/uname_change")
def uname_change(data: changeuname):
    user_id = data.user_id
    user_name = data.user_name
    
    connection = create_connection()
    cursor = connection.cursor()
    # 检查用户名是否已被其他用户使用
    cursor.execute("SELECT * FROM user WHERE user_name=%s AND user_id!=%s", (user_name, user_id))
    existing_user = cursor.fetchone()
    if existing_user:
        return {"success": False, "message": "用户名已存在"}
    else:
        # 检查用户是否存在
        cursor.execute("SELECT * FROM user WHERE user_id=%s", (user_id,))
        user = cursor.fetchone()
        if user:
            cursor.execute("UPDATE user SET user_name=%s WHERE user_id=%s", (user_name, user_id))
            connection.commit()
            return {"success": True, "message": "用户名修改成功"}
        else:
            return {"success": False, "message": "用户不存在"}


# 获取所有用户列表
@app.get("/users")
def get_users():
    connection = create_connection()
    try:
        # 检查连接是否有效
        if connection is None or not connection.open:
            connection = create_connection()
            if connection is None:
                return {"success": False, "message": "无法连接到数据库"}
        
        cursor = connection.cursor()
        cursor.execute("SELECT user_id as id, user_name, email, pwd as password FROM user ORDER BY user_id")
        users = cursor.fetchall()
        
        # 添加默认注册时间和角色信息
        for user in users:
            user['registerTime'] = '2024-01-01 00:00:00'
            # 设置角色：user_id为0的是管理员，其他是普通用户
            if user['id'] == 0:
                user['role'] = 'admin'
            else:
                user['role'] = 'user'
        
        return {"success": True, "users": users}
    except Exception as e:
        return {"success": False, "message": f"获取用户列表失败: {str(e)}"}


# 删除用户
@app.delete("/delete_user")
def delete_user(user_id: int):
    connection = create_connection()
    try:
        # 检查连接是否有效
        if connection is None or not connection.open:
            connection = create_connection()
            if connection is None:
                return {"success": False, "message": "无法连接到数据库"}
        
        cursor = connection.cursor()
        cursor.execute("DELETE FROM user WHERE user_id = %s", (user_id,))
        connection.commit()
        
        if cursor.rowcount > 0:
            return {"success": True, "message": "用户删除成功"}
        else:
            return {"success": False, "message": "用户不存在"}
    except Exception as e:
        return {"success": False, "message": f"删除用户失败: {str(e)}"}

class update_user(BaseModel):
    data: dict

# 更新用户信息
@app.post("/update_user")
def update_user(data: update_user):
    user_id = data.data["user_id"]
    user_name = data.data["user_name"]
    email = data.data["email"]
    
    connection = create_connection()
    try:
        # 检查连接是否有效
        if connection is None or not connection.open:
            connection = create_connection()
            if connection is None:
                return {"success": False, "message": "无法连接到数据库"}
        
        cursor = connection.cursor()
        
        # 构建更新语句
        update_fields = []
        update_values = []
        
        if user_name is not None:
            update_fields.append("user_name = %s")
            update_values.append(user_name)
        
        if email is not None:
            update_fields.append("email = %s")
            update_values.append(email)
        
        if not update_fields:
            return {"success": False, "message": "没有提供更新字段"}
        
        update_values.append(user_id)
        update_query = f"UPDATE user SET {', '.join(update_fields)} WHERE user_id = %s"
        
        cursor.execute(update_query, update_values)
        connection.commit()
        
        if cursor.rowcount > 0:
            return {"success": True, "message": "用户信息更新成功"}
        else:
            return {"success": False, "message": "用户不存在或信息未改变"}
    except Exception as e:
        return {"success": False, "message": f"更新用户信息失败: {str(e)}"}


# 重命名主题
@app.get("/subject_change")
def subject_change(subject_id: int = Query(...), new_subject: str = Query(...)):
    
    try:
        connection = create_connection()
        if not connection:
            return {"success": False, "message": "数据库连接失败"}

        cursor = connection.cursor()
        
        # 检查主题是否存在
        cursor.execute("SELECT * FROM subject WHERE subject_id = %s", (subject_id,))
        result = cursor.fetchone()
        if not result:
            return {"success": False, "message": "主题不存在"}

        # 更新主题名称
        cursor.execute("UPDATE subject SET subject_name = %s WHERE subject_id = %s", 
                      (new_subject, subject_id))
        connection.commit()

        if cursor.rowcount == 0:
            return {"success": False, "message": "未更新任何记录"}

        return {"success": True, "message": "主题名称已更新"}

    except Exception as e:
        if connection:
            connection.rollback()
        return {"success": False, "message": f"数据库错误: {str(e)}"}
    finally:
        cursor.close()
        

    

# 问答接口（流式输出）
@app.get("/stream_chat")
async def stream_chat(user_id: int, question: str, subjectid: int):
    # 创建新的数据库连接
    db_conn = create_connection()
    if db_conn is None:
        return "错误: 无法连接到数据库"
        
    try:
        cursor = db_conn.cursor()
        
        # 先保存用户问题到数据库（确保用户问题能立即显示）
        try:
            # 如果是新主题，先创建一个临时主题ID用于保存问题
            temp_subjectid = subjectid
            if subjectid == 0:
                # 先插入一个临时主题
                cursor.execute("INSERT INTO subject (subject_name) VALUES (%s)", ("临时主题",))
                db_conn.commit()
                temp_subjectid = cursor.lastrowid
                print(f"创建临时主题，ID: {temp_subjectid}")
                
                # 验证主题是否真的创建成功
                cursor.execute("SELECT * FROM subject WHERE subject_id = %s", (temp_subjectid,))
                subject_check = cursor.fetchone()
                if not subject_check:
                    print(f"警告: 主题创建失败，主题ID {temp_subjectid} 不存在")
                else:
                    print(f"主题创建验证成功: {subject_check}")
            
            cursor.execute("INSERT INTO chatcontent (user_id, subject_id, question) VALUES (%s, %s, %s)", 
                          (user_id, temp_subjectid, question))
            db_conn.commit()
            chat_id = cursor.lastrowid
            print(f"已保存用户问题，主题ID: {temp_subjectid}")
        except Exception as e:
            print(f"保存用户问题失败: {str(e)}")
        
        # 处理主题ID为0的情况（在后台生成真实主题）
        final_subjectid = temp_subjectid
        if subjectid == 0:
            try:
                # 调用LLM生成主题
                subject = llm(question)
                # 更新主题名称
                cursor.execute("UPDATE subject SET subject_name = %s WHERE subject_id = %s", (subject, final_subjectid))
                db_conn.commit()
                print(f"更新主题: {subject}, 主题ID: {final_subjectid}")
            except Exception as e:
                print(f"生成主题失败: {str(e)}")
        
        # 定义异步生成器用于流式输出
        async def generate_response():
            try:
                # 如果是新创建的主题，先返回主题ID
                if subjectid == 0:
                    yield f"{final_subjectid}"
                
                # 直接开始生成回答，不输出主题ID
                result = llm_stream(question, final_subjectid)
                full_content = ""
                
                for el in result:
                    if el.startswith("错误:"):
                        # 如果返回的是错误信息，直接返回
                        yield el
                        return
                    # 流式输出每个内容块
                    yield el
                    full_content += el
                
                # 保存完整AI回答到数据库
                if full_content:
                    try:
                        new_conn = create_connection()  # 创建新连接避免冲突
                        if new_conn:
                            new_cursor = new_conn.cursor()
                            new_cursor.execute("UPDATE chatcontent SET  answer = %s WHERE chat_id = %s AND subject_id = %s", 
                                              (full_content, chat_id ,final_subjectid))
                            new_conn.commit()
                            print(f"已保存AI回答，主题ID: {final_subjectid}")
                            new_cursor.close()
                            new_conn.close()
                    except Exception as e:
                        print(f"保存AI回答失败: {str(e)}")
            except Exception as e:
                error_msg = f"错误: 生成回答失败: {str(e)}"
                print(error_msg)
                yield error_msg

        # 使用StreamingResponse实现真正的流式输出，设置长超时时间
        return StreamingResponse(
            generate_response(), 
            media_type="text/plain",
            headers={
                "X-Accel-Buffering": "no",  # 禁用Nginx缓冲
                "Cache-Control": "no-cache"  # 禁用缓存
            }
        )
    except Exception as e:
        error_msg = f"错误: 处理请求时出错: {str(e)}"
        print(error_msg)
        return error_msg
    finally:
        # 关闭数据库连接
        try:
            if cursor:
                cursor.close()
            if db_conn:
                db_conn.close()
        except Exception as e:
            print(f"关闭数据库连接失败: {str(e)}")



# 获取用户的历史对话主题列表
@app.get("/get_subjects")
def get_subjects(user_id: int):
    """获取用户的所有历史对话主题"""
    db_conn = create_connection()
    if db_conn is None:
        return JSONResponse(content={"success": False, "message": "无法连接到数据库"})
    
    try:
        cursor = db_conn.cursor()
        # 查询用户的所有主题
        cursor.execute("""
            SELECT DISTINCT s.subject_id, s.subject_name, s.created_at
            FROM subject s
            JOIN chatcontent c ON s.subject_id = c.subject_id
            WHERE c.user_id = %s
            ORDER BY s.created_at DESC
        """, (user_id,))
        subjects = cursor.fetchall()
        
        # 确保返回的数据编码正确并处理datetime对象
        processed_subjects = []
        for subject in subjects:
            subject_dict = dict(subject)
            
            # 处理主题名称编码
            if isinstance(subject_dict['subject_name'], bytes):
                subject_dict['subject_name'] = subject_dict['subject_name'].decode('utf-8')
            
            # 将datetime对象转换为字符串
            if isinstance(subject_dict['created_at'], datetime):
                subject_dict['created_at'] = subject_dict['created_at'].isoformat()
            
            processed_subjects.append(subject_dict)
        
        # 使用ensure_ascii=False确保中文正确显示
        import json
        json_str = json.dumps({"success": True, "subjects": processed_subjects}, ensure_ascii=False)
        return JSONResponse(content=json.loads(json_str))
    except Exception as e:
        return JSONResponse(content={"success": False, "message": f"获取历史对话失败: {str(e)}"})
    finally:
        if db_conn and db_conn.open:
            db_conn.close()

# 获取特定主题的对话内容
@app.get("/get_chat_history")
def get_chat_history(user_id: int, subject_id: int):
    """获取特定主题的完整对话历史"""
    db_conn = create_connection()
    if db_conn is None:
        return {"success": False, "message": "无法连接到数据库"}
    
    try:
        cursor = db_conn.cursor()
        # 查询该主题的所有对话内容
        cursor.execute("""
            SELECT 
                chat_id,
                user_id,
                subject_id,
                question,
                answer,
                created_at
            FROM chatcontent
            WHERE user_id = %s AND subject_id = %s
            ORDER BY created_at ASC
        """, (user_id, subject_id))
        chat_history = cursor.fetchall()
        
        # 确保返回的数据格式正确
        formatted_history = []
        for item in chat_history:
            # 处理字典格式的查询结果
            if isinstance(item, dict):
                formatted_history.append({
                    "chat_id": item["chat_id"],
                    "user_id": item["user_id"],
                    "subject_id": item["subject_id"],
                    "question": item["question"],
                    "answer": item["answer"],
                    "created_at": item["created_at"].strftime("%Y-%m-%d %H:%M:%S") if item["created_at"] else None
                })
            else:
                # 处理元组格式的查询结果
                formatted_history.append({
                    "chat_id": item[0],
                    "user_id": item[1],
                    "subject_id": item[2],
                    "question": item[3],
                    "answer": item[4],
                    "created_at": item[5].strftime("%Y-%m-%d %H:%M:%S") if item[5] else None
                })
        
        print(f"返回聊天历史数据: 用户{user_id}, 主题{subject_id}, 记录数{len(formatted_history)}")
        
        return {"success": True, "chat_history": formatted_history}
    except Exception as e:
        print(f"获取对话历史失败: {str(e)}")
        return {"success": False, "message": f"获取对话历史失败: {str(e)}"}
    finally:
        if db_conn and db_conn.open:
            db_conn.close()

# 删除历史对话主题及所有对话内容
@app.delete("/delete_subject")
def delete_subject(user_id: int, subject_id: int):
    """删除指定主题及其所有对话内容"""
    db_conn = create_connection()
    if db_conn is None:
        return {"success": False, "message": "无法连接到数据库"}
    
    try:
        cursor = db_conn.cursor()
        
        # 首先删除该主题下的所有对话内容
        cursor.execute("DELETE FROM chatcontent WHERE user_id = %s AND subject_id = %s", (user_id, subject_id))
        
        # 然后删除主题本身
        cursor.execute("DELETE FROM subject WHERE subject_id = %s", (subject_id,))
        
        db_conn.commit()
        
        return {"success": True, "message": "历史对话删除成功"}
    except Exception as e:
        db_conn.rollback()
        print(f"删除历史对话失败: {str(e)}")
        return {"success": False, "message": f"删除历史对话失败: {str(e)}"}
    finally:
        if db_conn and db_conn.open:
            db_conn.close()



# 服务器启动入口
if __name__ == "__main__":
    # 检查依赖
    if not check_dependencies():
        sys.exit(1)
    
    # 确保static目录存在
    ensure_static_dir()
    
    # 配置服务器参数
    host = "0.0.0.0"  # 监听所有网络接口
    port = 8000      # 默认端口
    reload = True    # 开发模式下自动重载
    
    # 导入uvicorn并启动服务器
    try:
        import uvicorn
        uvicorn.run(
            "server:app",
            host=host,
            port=port,
            reload=reload,
            reload_dirs=[".", "servers/"]  # 修复为目录列表
        )
    except ImportError:
        sys.exit(1)
    except Exception as e:
        sys.exit(1)
        
    


        
    